Repository navigation
Optimize FP16 LayerNorm and RMSNorm on x86 - #32715
Conversation
Add an AVX2/F16C MLAS kernel for FP16 LayerNorm and RMSNorm and dispatch it from the CPU implementation while retaining the scalar fallback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The kernel breaks 32-bit Unix builds without -mf16c and returns an incorrect simplified MeanOut value.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
Open (3)
What changed in this PR
Adds an AVX2/FMA/F16C FP16 LayerNorm/RMSNorm kernel for x86 and dispatches eligible CPU workloads to it.
Changes:
- Implements vectorized FP16 normalization with FP32 reductions.
- Adds CPU feature dispatch and MLAS API plumbing.
- Adds MLAS and CPU-provider tests.
| File | Description |
|---|---|
onnxruntime/test/providers/cpu/nn/rms_norm_op_test.cc |
Tests an optimized FP16 RMSNorm row. |
onnxruntime/test/mlas/unittest/test_layernorm.cpp |
Tests FP16 MLAS normalization. |
onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc |
Dispatches eligible FP16 rows to MLAS. |
onnxruntime/core/mlas/lib/platform.cpp |
Registers the kernel on supported x86 CPUs. |
onnxruntime/core/mlas/lib/mlasi.h |
Defines internal kernel types and dispatch state. |
onnxruntime/core/mlas/lib/layernorm.cpp |
Adds the FP16 dispatch wrapper. |
onnxruntime/core/mlas/lib/layernorm_kernel_avx2.cpp |
Implements the AVX2 FP16 kernel. |
onnxruntime/core/mlas/inc/mlas.h |
Exposes the FP16 MLAS API. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Enable F16C when compiling the 32-bit x86 kernel, return the RMSNorm mean when requested, and cover optimized LayerNorm provider wiring end to end. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Use the supported round-to-nearest immediate for F16C conversions. MSVC rejects _MM_FROUND_NO_EXC with C4556 because only bits 0-2 are valid; the extra bit is ignored by VCVTPS2PH and does not suppress exceptions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Fixed the Windows MSVC build failure in ba7dfc0. The three FP16 output conversions passed The three existing review threads are already addressed and resolved. The push starts fresh CI for this fix; old failed runs were not rerun. |
Description
Add an AVX2/FMA/F16C MLAS kernel for FP16 LayerNorm and RMSNorm on x86. The kernel widens FP16 inputs in vectors, performs the reductions and affine transform in FP32 (with a double-precision mean accumulation for LayerNorm), and narrows the result back to FP16 in vectors.
The CPU LayerNorm implementation dispatches to this kernel for rows of at least 16 elements when the processor supports the required instructions. Other architectures, older x86 processors, short rows, BF16, and unsupported output-statistics types retain the existing scalar path.
Motivation and Context
The existing FP16 CPU path converts every input and output element individually around scalar normalization math. Transformer workloads commonly execute LayerNorm and RMSNorm over rows of 1K-4K elements, where conversion and scalar arithmetic dominate execution time.
The benchmark processes 4,194,304 FP16 elements per invocation with rows parallelized across OpenMP workers, mirroring ORT's per-row thread-pool parallelism. It was pinned to one 48-core Intel Xeon Platinum 8480C socket with workers bound to cores. The table reports the median speedup from three process runs; graph/session overhead is intentionally excluded.
At 32 threads, the optimized kernels reach median throughputs of 54.0-57.6 billion elements/s for LayerNorm and 81.2-89.4 billion elements/s for RMSNorm.
Validation:
onnxruntime_mlas_testandonnxruntime_provider_testin a CPU Release configuration.MlasLayerNormF16Test.MatchesFloatReferencepasses for LayerNorm with/without bias and RMSNorm at row sizes 16, 31, and 128.RMSNormalizationOpTest.RMSNorm_float16_OptimizedRowpasses through the CPU execution provider.RMSNormalizationOpTest.*suite passes (26 tests before adding the dedicated optimized-row case).